package com.zenithmind.user.security;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.Map;
import java.util.UUID;

/**
 * OAuth2 Token Utilities for generating JWT tokens
 */
@Component
public class OAuth2TokenUtils {

    private final JwtEncoder jwtEncoder;
    private final JwtDecoder jwtDecoder;
    
    @Value("${spring.security.oauth2.resourceserver.jwt.issuer-uri:http://localhost:8081}")
    private String issuer;
    
    @Value("${jwt.access-token-expiration:3600}")
    private int accessTokenExpiration = 3600; // Default 1 hour

    @Autowired
    public OAuth2TokenUtils(JwtEncoder jwtEncoder, JwtDecoder jwtDecoder) {
        this.jwtEncoder = jwtEncoder;
        this.jwtDecoder = jwtDecoder;
    }
    
    /**
     * Create a JWT token with the provided claims
     *
     * @param claims Claims to include in the token
     * @return JWT token string
     */
    public String createJwt(Map<String, Object> claims) {
        Instant issuedAt = Instant.now();
        Instant expiresAt = issuedAt.plus(accessTokenExpiration, ChronoUnit.SECONDS);
        
        JwtClaimsSet.Builder claimsBuilder = JwtClaimsSet.builder()
                .issuer(issuer)
                .issuedAt(issuedAt)
                .expiresAt(expiresAt)
                .id(UUID.randomUUID().toString());
        
        // Add standard sub claim if available
        if (claims.containsKey("sub")) {
            claimsBuilder.subject((String) claims.get("sub"));
            claims.remove("sub"); // Remove to prevent duplication
        }
        
        // Add all other provided claims
        claims.forEach(claimsBuilder::claim);
        
        // Encode and create JWT
        return jwtEncoder.encode(JwtEncoderParameters.from(claimsBuilder.build())).getTokenValue();
    }
    
    /**
     * Validate a JWT token and return its claims
     *
     * @param token The JWT token string to validate
     * @return The decoded JWT object with claims if valid
     */
    public Jwt validateToken(String token) {
        return jwtDecoder.decode(token);
    }
    
    /**
     * Extract user ID from JWT token
     *
     * @param token JWT token
     * @return User ID if present, null otherwise
     */
    public String getUserIdFromToken(String token) {
        Jwt jwt = validateToken(token);
        return jwt.getClaim("user_id");
    }
} 